import torch
from easycore.common.registry import Registry


class OptimizerFactory:
    registry = Registry("optimizer")

    @classmethod
    def register(cls, optimizer_type_name=None, obj=None):
        return cls.registry.register(optimizer_type_name, obj)

    @classmethod
    def get(cls, optimizer_config, model):
        optimizer = cls.registry.get(optimizer_config['type']).build(optimizer_config, model)
        return optimizer


@OptimizerFactory.register()
class AdamW:
    @classmethod
    def build(cls, optimizer_config, model):
        optimizer = torch.optim.AdamW(model.parameters(), optimizer_config['learning_rate'])
        return optimizer


@OptimizerFactory.register()
class SGD:
    @classmethod
    def build(cls, optimizer_config, model):
        optimizer = torch.optim.SGD(model.parameters(), optimizer_config['learning_rate'], momentum=optimizer_config['momentum'], nesterov=True)
        return optimizer
